#include <kuai/sim/energy.hpp>
#include <kuai/sim/efunctor.hpp>

#define TRY_TO_CREATE(function) if (boost::algorithm::ends_with(name,		\
		boost::algorithm::to_lower_copy(String(#function))))				\
	{																		\
		if (pbc == NULL) {													\
			if (name == "bond/"												\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionBond<Functor##function, false>;			\
			}																\
			else if (name == "angle/cos_"									\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionCosAngle<Functor##function, false>;		\
			}																\
			else if (name == "angle/"										\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionAngle<Functor##function, false>;			\
			}																\
			else if (name == "torsion/cos_"									\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionCosTorsion<Functor##function, false>;	\
			}																\
			else if (name == "torsion/"										\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionTorsion<Functor##function, false>;		\
			}																\
			else if (name == "improper_torsion/cos_"						\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionCosTorsion<Functor##function, false>;	\
			}																\
			else if (name == "improper_torsion/"							\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionTorsion<Functor##function, false>;		\
			}																\
		}																	\
		else {																\
			if (name == "bond/"												\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionBond<Functor##function, true>;			\
			}																\
			else if (name == "angle/cos_"									\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionCosAngle<Functor##function, true>;		\
			}																\
			else if (name == "angle/"										\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionAngle<Functor##function, true>;			\
			}																\
			else if (name == "torsion/cos_"									\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionCosTorsion<Functor##function, true>;		\
			}																\
			else if (name == "torsion/"										\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionTorsion<Functor##function, true>;		\
			}																\
			else if (name == "improper_torsion/cos_"						\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionCosTorsion<Functor##function, true>;		\
			}																\
			else if (name == "improper_torsion/"							\
				+ boost::algorithm::to_lower_copy(String( #function ))) {	\
				return new FunctionTorsion<Functor##function, true>;		\
			}																\
		}																	\
	}																		\


namespace kuai {

	EnergySet::EnergySet() {
		nAtoms = 0;
	}

	void EnergySet::start(SimulationEnvironment& env) {
		if (env.getTID() == 0) {
			rank = getRank();
			eBuffer.resize(rank * (env.getCPU() - 1));
			gBuffer.resize(nAtoms * (env.getCPU() - 1));
		}
		env.wait();
		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->start(env);
		}
	}
		
	void EnergySet::addE(SimulationEnvironment& env, RealNumber e[]) {
		assert (env.getCPU() > 0);
		assert (eBuffer.size() == rank * (env.getCPU()-1));

		SizePair size = env.range(eBuffer.size());
		memset(&eBuffer[size.first], 0, sizeof(RealNumber) * (size.second-size.first));

		RealNumber* result;
		if (env.getTID() == 0) {
			result = e;
		}
		else {
			result = &eBuffer[0] + rank*(env.getTID()-1);
		}
		env.wait();

		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->addE(env, result);
			result += (*it)->rank;
		}
		env.wait();

		size = env.range(rank);
		for (size_t i = size.first; i < size.second; ++i) {
			for (size_t j = 0; j < env.getCPU()-1; ++j)
			{
				e[i] += eBuffer[j*rank + i]; 
			}
		}
		env.wait();
	}
		
	void EnergySet::addEG(SimulationEnvironment& env, RealNumber e[], XYZ gred[]) {
		assert (env.getCPU() > 0);
		assert (eBuffer.size() == rank * (env.getCPU()-1));
		assert (gBuffer.size() == nAtoms * (env.getCPU()-1));

		SizePair size = env.range(eBuffer.size());
		memset(&eBuffer[size.first], 0, sizeof(RealNumber) * (size.second-size.first));
		size = env.range(gBuffer.size());
		memset(&gBuffer[size.first], 0, sizeof(XYZ) * (size.second-size.first));

		RealNumber* eresult;
		XYZ* gresult;
		if (env.getTID() == 0) {
			eresult = e;
			gresult = gred;
		}
		else {
			eresult = &eBuffer[rank*(env.getTID()-1)];
			gresult = &gBuffer[nAtoms*(env.getTID()-1)];
		}
		env.wait();

		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->addEG(env, eresult, gresult);
			eresult += (*it)->rank;
		}
		env.wait();

		size = env.range(rank);
		for (size_t j = 0; j < env.getCPU()-1; ++j) {
			for (size_t i = size.first; i < size.second; ++i) {
				e[i] += eBuffer[j*rank + i]; 
			}
		}
		size = env.range(nAtoms);
		for (size_t j = 0; j < env.getCPU()-1; ++j) {
			for (size_t i = size.first; i < size.second; ++i) {
				gred[i] += gBuffer[j*nAtoms + i]; 
			}
		}
		env.wait();
	}

	void EnergySet::addHX(SimulationEnvironment& env, const RealNumber x[], RealNumber result[]) {
		assert (env.getCPU() > 0);
		assert (gBuffer.size() == nAtoms * (env.getCPU()-1));

		SizePair size = size = env.range(gBuffer.size());
		memset(&gBuffer[size.first], 0, sizeof(XYZ) * (size.second-size.first));

		RealNumber* buffer;
		if (env.getTID() == 0) {
			buffer = result;
		}
		else {
			buffer = reinterpret_cast<RealNumber*>(&gBuffer[nAtoms*(env.getTID()-1)]);
		}
		env.wait();

		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->addHX(env, x, buffer);
		}
		env.wait();

		size = env.range(nAtoms);
		XYZ* pResult = reinterpret_cast<XYZ*>(result);
		for (size_t j = 0; j < env.getCPU()-1; ++j) {
			for (size_t i = size.first; i < size.second; ++i) {
				pResult[i] += gBuffer[j*nAtoms + i]; 
			}
		}
		env.wait();
	}


	void EnergySet::addEGH(SimulationEnvironment& env, RealNumber e[], XYZ gred[], RealNumber* hess[]) {
		assert (env.getCPU() > 0);
		assert (eBuffer.size() == rank * (env.getCPU()-1));
		assert (gBuffer.size() == nAtoms * (env.getCPU()-1));

		SizePair size = env.range(eBuffer.size());
		memset(&eBuffer[size.first], 0, sizeof(RealNumber) * (size.second-size.first));
		size = env.range(gBuffer.size());
		memset(&gBuffer[size.first], 0, sizeof(XYZ) * (size.second-size.first));

		RealNumber* eresult;
		XYZ* gresult;
		if (env.getTID() == 0) {
			eresult = e;
			gresult = gred;
		}
		else {
			eresult = &eBuffer[rank*(env.getTID()-1)];
			gresult = &gBuffer[nAtoms*(env.getTID()-1)];
		}
		env.wait();

		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->addEGH(env, eresult, gresult, hess);
			eresult += (*it)->rank;
		}
		env.wait();

		size = env.range(rank);
		for (size_t j = 0; j < env.getCPU()-1; ++j) {
			for (size_t i = size.first; i < size.second; ++i) {
				e[i] += eBuffer[j*rank + i]; 
			}
		}
		size = env.range(nAtoms);
		for (size_t j = 0; j < env.getCPU()-1; ++j) {
			for (size_t i = size.first; i < size.second; ++i) {
				gred[i] += gBuffer[j*nAtoms + i]; 
			}
		}
		env.wait();
	}
		
	void EnergySet::getE(SimulationEnvironment& env, RealNumber e[]) {
		std::pair<size_t, size_t> range = env.range(rank);
		memset(e+range.first, 0, sizeof(RealNumber)*(range.second-range.first));
		env.wait();
		addE(env, e);
	}
	
	void EnergySet::getEG(SimulationEnvironment& env, RealNumber e[], XYZ gred[]) {
		std::pair<size_t, size_t> range = env.range(rank);
		memset(e+range.first, 0, sizeof(RealNumber)*(range.second-range.first));
		range = env.range(nAtoms);
		memset(gred+range.first, 0, sizeof(XYZ)*(range.second -range.first));
		env.wait();
		addEG(env, e, gred);
	}
	void EnergySet::getEGH(SimulationEnvironment& env, RealNumber e[], XYZ gred[], RealNumber* hess[]) {
		std::pair<size_t, size_t> range = env.range(rank);
		memset(e+range.first, 0, sizeof(RealNumber)*(range.second-range.first));
		range = env.range(nAtoms);
		memset(gred+range.first, 0, sizeof(XYZ)*(range.second -range.first));
		for (size_t i = range.first*3; i < range.second*3; ++i) {
			memset(hess[i], 0, sizeof(RealNumber) * 3 * nAtoms);
		}
		env.wait();
		addEGH(env, e, gred, hess);
	}
	void EnergySet::getHX(SimulationEnvironment& env, const RealNumber x[], RealNumber result[]) {
		std::pair<size_t, size_t> range = range = env.range(3*nAtoms);
		memset(result+range.first, 0, sizeof(RealNumber)*(range.second -range.first));
		env.wait();
		addHX(env, x, result);
	}
	
	void EnergySet::finish(SimulationEnvironment& env) {
		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->finish(env);
		}
		eBuffer.clear();
		gBuffer.clear();
	}
		
	void EnergySet::finalize(SimulationEnvironment& env) {
		for (std::vector<EnergyFunction*>::iterator
			it = funcs.begin(); it != funcs.end(); ++it)
		{
			(*it)->finalize(env);
		}
		eBuffer.clear();
		gBuffer.clear();
	}




	BasicForceFieldEnergy* createValenceEnergyFunction(const char szName[], const PBC* pbc) {
		const String name = boost::algorithm::to_lower_copy(String(szName));
		
		TRY_TO_CREATE(Harmonic);
		TRY_TO_CREATE(Quardric);
		TRY_TO_CREATE(Poly5);
		TRY_TO_CREATE(LJ12_6);
		TRY_TO_CREATE(LJ_9_6);
		
		return NULL;
	}
	
}
